package com.mkx.orderedlist.skiplist;

import com.mkx.orderedlist.OrderedList;
import java.util.Random;

/**
 * 基于跳表实现的有序表
 * @param <K> 键类型
 * @param <V> 值类型
 */
public class SkipListOrderedList<K extends Comparable<K>, V> implements OrderedList<K, V> {
    private static final int MAX_LEVEL = 16; // 最大层数
    private final Random random;
    private int level; // 当前跳表的最高层数
    private final SkipListNode<K, V> header; // 头节点（哨兵节点）
    private int size; // 元素个数
    
    /**
     * 构造函数
     */
    @SuppressWarnings("unchecked")
    public SkipListOrderedList() {
        this.random = new Random();
        this.level = 0;
        this.size = 0;
        // 创建头节点，层数为MAX_LEVEL
        this.header = new SkipListNode<>(null, null, MAX_LEVEL);
    }
    
    /**
     * 随机生成节点的层数
     * 遵循概率分布：有50%的概率增加一层
     * @return 生成的层数
     */
    private int randomLevel() {
        int level = 0;
        // 50%的概率增加层数，但不超过最大层数
        while (random.nextBoolean() && level < MAX_LEVEL) {
            level++;
        }
        return level;
    }
    
    /**
     * 查找键对应的节点
     * @param key 键
     * @return 对应的节点，如果不存在返回null
     */
    private SkipListNode<K, V> findNode(K key) {
        SkipListNode<K, V> current = header;
        
        // 从最高层开始查找
        for (int i = level; i >= 0; i--) {
            // 在当前层向前移动，直到下一个节点的键大于等于目标键
            while (current.forward[i] != null && current.forward[i].key.compareTo(key) < 0) {
                current = current.forward[i];
            }
        }
        
        // 下降到第0层，检查下一个节点是否就是目标键
        current = current.forward[0];
        
        // 验证是否找到
        if (current != null && current.key.compareTo(key) == 0) {
            return current;
        }
        
        return null;
    }
    
    /**
     * 查找键，并记录查找路径（用于插入和删除操作）
     * @param key 键
     * @param update 用于存储查找路径上的节点
     * @return 目标节点的前一个节点在第0层的位置
     */
    @SuppressWarnings("unchecked")
    private SkipListNode<K, V> findWithUpdate(K key, SkipListNode<K, V>[] update) {
        SkipListNode<K, V> current = header;
        
        // 从最高层开始查找
        for (int i = level; i >= 0; i--) {
            while (current.forward[i] != null && current.forward[i].key.compareTo(key) < 0) {
                current = current.forward[i];
            }
            update[i] = current;
        }
        
        return current.forward[0]; // 返回第0层的下一个节点
    }
    
    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        @SuppressWarnings("unchecked")
        SkipListNode<K, V>[] update = new SkipListNode[MAX_LEVEL + 1];
        
        // 查找插入位置并记录路径
        SkipListNode<K, V> current = findWithUpdate(key, update);
        
        // 如果键已存在，更新值
        if (current != null && current.key.compareTo(key) == 0) {
            V oldValue = current.value;
            current.value = value;
            return oldValue;
        }
        
        // 生成新节点的层数
        int newLevel = randomLevel();
        
        // 如果新节点的层数高于当前跳表的最高层数，更新最高层数和update数组
        if (newLevel > level) {
            for (int i = level + 1; i <= newLevel; i++) {
                update[i] = header;
            }
            level = newLevel;
        }
        
        // 创建新节点
        SkipListNode<K, V> newNode = new SkipListNode<>(key, value, newLevel);
        
        // 更新指针，插入新节点
        for (int i = 0; i <= newLevel; i++) {
            newNode.forward[i] = update[i].forward[i];
            update[i].forward[i] = newNode;
        }
        
        size++;
        return null;
    }
    
    @Override
    public V get(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        SkipListNode<K, V> node = findNode(key);
        return (node != null) ? node.value : null;
    }
    
    @Override
    public V remove(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        @SuppressWarnings("unchecked")
        SkipListNode<K, V>[] update = new SkipListNode[MAX_LEVEL + 1];
        
        // 查找要删除的节点并记录路径
        SkipListNode<K, V> current = findWithUpdate(key, update);
        
        // 如果节点不存在，返回null
        if (current == null || current.key.compareTo(key) != 0) {
            return null;
        }
        
        V oldValue = current.value;
        
        // 更新指针，删除节点
        for (int i = 0; i <= level; i++) {
            if (update[i].forward[i] != current) {
                break;
            }
            update[i].forward[i] = current.forward[i];
        }
        
        // 更新跳表的最高层数
        while (level > 0 && header.forward[level] == null) {
            level--;
        }
        
        size--;
        return oldValue;
    }
    
    @Override
    public boolean containsKey(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        return findNode(key) != null;
    }
    
    @Override
    public int size() {
        return size;
    }
    
    @Override
    public boolean isEmpty() {
        return size == 0;
    }
    
    @Override
    public void clear() {
        // 重置头节点的所有forward指针
        for (int i = 0; i <= level; i++) {
            header.forward[i] = null;
        }
        level = 0;
        size = 0;
    }
    
    @Override
    public K firstKey() {
        if (isEmpty()) {
            return null;
        }
        // 第0层的第一个有效节点就是最小键
        return header.forward[0].key;
    }
    
    @Override
    public K lastKey() {
        if (isEmpty()) {
            return null;
        }
        
        SkipListNode<K, V> current = header;
        // 从最高层开始，一直向右走
        for (int i = level; i >= 0; i--) {
            while (current.forward[i] != null) {
                current = current.forward[i];
            }
        }
        
        return current.key;
    }
}